package com.yjzx.llama.records;

import java.io.PrintStream;
import java.nio.file.Path;
import java.nio.file.Paths;

/**
 * @param modelPath    模型路径
 * @param prompt       单次模式下提问信息
 * @param systemPrompt 系统提示符
 * @param interactive  是否聊天模式
 * @param temperature
 * @param topp
 * @param seed         随机种子
 * @param maxTokens    最大的tokens长度
 * @param stream       是否生成过程中打印令牌
 * @param echo         将 ALL 令牌打印到 stderr
 */
public record Options(Path modelPath, String prompt, String systemPrompt, boolean interactive,
                      float temperature, float topp, long seed, int maxTokens, boolean stream, boolean echo) {

    public static final int DEFAULT_MAX_TOKENS = 512;

    public Options {
        require(modelPath != null, "缺少参数: --model <路径> 是必须的");
        require(interactive || prompt != null, "缺少参数：--prompt 在 --instruct 模式下是必需的，例如 --prompt \"Why is the sky blue？\"");
        require(0 <= temperature, "无效参数: --temperature 必须为非负数");
        require(0 <= topp && topp <= 1, "无效参数: --top-p 必须在 [0, 1]");
    }

    public static void require(boolean condition, String messageFormat, Object... args) {
        if (!condition) {
            System.out.println("ERROR " + messageFormat.formatted(args));
            System.out.println();
            printUsage(System.out);
            System.exit(-1);
        }
    }

    static void printUsage(PrintStream out) {
        out.println("Usage:  jbang Llama3.java [options]");
        out.println();
        out.println("Options:");
        out.println("  --model, -m <path>            必需，.gguf 文件的路径");
        out.println("  --interactive, --chat, -i     在聊天模式下运行");
        out.println("  --instruct                    在 Instruct （单次执行，配合输入提示） 模式下运行，默认模式");
        out.println("  --prompt, -p <string>         输入提示");
        out.println("  --system-prompt, -sp <string> (optional) 系统提示符");
        out.println("  --temperature, -temp <float>  temperature in [0,inf], default 0.1");
        out.println("  --top-p <float>               p 值 in top-p (nucleus) 采样在 [0,1] 默认 0.95");
        out.println("  --seed <long>                 随机种子, 默认 System.nanoTime()");
        out.println("  --max-tokens, -n <int>        上下文长度, 默认值 " + DEFAULT_MAX_TOKENS);
        out.println("  --stream <boolean>            在生成过程中打印令牌;可能会导致非 ASCII 文本的编码乱码，默认为 true");
        out.println("  --echo <boolean>              将 ALL 令牌打印到 stderr，如果为 true，建议设置 --stream=false，默认为 false");
        out.println();
        out.println("Examples:");
        out.println("  jbang Llama3.java --model llama3.2-1b-q4_0.gguf --prompt \"Tell me a joke\"");
        out.println("  jbang Llama3.java --model llama3.2-1b-q4_0.gguf --system-prompt \"Reply concisely, in French\" --prompt \"Who was Marie Curie?\"");
        out.println("  jbang Llama3.java --model llama3.2-1b-q4_0.gguf --system-prompt \"Answer concisely\" --chat");
        out.println("  jbang Llama3.java --model llama3.2-1b-q4_0.gguf --chat");
        out.println("  jbang Llama3.java --model llama3.2-1b-q4_0.gguf --prompt \"Print 5 emojis\" --stream=false");
    }

    public static Options parseOptions(String[] args) {
        String prompt = null;
        String systemPrompt = null;
        float temperature = 0.1f;
        float topp = 0.95f;
        Path modelPath = null;
        long seed = System.nanoTime();
        // Keep max context length small for low-memory devices.
        int maxTokens = DEFAULT_MAX_TOKENS;
        boolean interactive = false;
        boolean stream = true;
        boolean echo = false;

        for (int i = 0; i < args.length; i++) {
            String optionName = args[i];
            require(optionName.startsWith("-"), "Invalid option %s", optionName);
            switch (optionName) {
                case "--interactive", "--chat", "-i" -> interactive = true;
                case "--instruct" -> interactive = false;
                case "--help", "-h" -> {
                    printUsage(System.out);
                    System.exit(0);
                }
                default -> {
                    String nextArg;
                    if (optionName.contains("=")) {
                        String[] parts = optionName.split("=", 2);
                        optionName = parts[0];
                        nextArg = parts[1];
                    } else {
                        require(i + 1 < args.length, "Missing argument for option %s", optionName);
                        nextArg = args[i + 1];
                        i += 1; // skip arg
                    }
                    switch (optionName) {
                        case "--prompt", "-p" -> prompt = nextArg;
                        case "--system-prompt", "-sp" -> systemPrompt = nextArg;
                        case "--temperature", "--temp" -> temperature = Float.parseFloat(nextArg);
                        case "--top-p" -> topp = Float.parseFloat(nextArg);
                        case "--model", "-m" -> modelPath = Paths.get(nextArg);
                        case "--seed", "-s" -> seed = Long.parseLong(nextArg);
                        case "--max-tokens", "-n" -> maxTokens = Integer.parseInt(nextArg);
                        case "--stream" -> stream = Boolean.parseBoolean(nextArg);
                        case "--echo" -> echo = Boolean.parseBoolean(nextArg);
                        default -> require(false, "Unknown option: %s", optionName);
                    }
                }
            }
        }
        return new Options(modelPath, prompt, systemPrompt, interactive, temperature, topp, seed, maxTokens, stream, echo);
    }
}
